import networkx as nx
import matplotlib.pyplot as plt
import random
import json
from networkx.readwrite import json_graph


def generate_random_topology(num_nodes, edge_probability):
    """
    生成随机网络拓扑
    :param num_nodes: 节点数量
    :param edge_probability: 边的生成概率
    :return: 生成的图对象
    """
    # 生成随机图
    G = nx.erdos_renyi_graph(num_nodes, edge_probability)

    # 为每个节点添加特征
    for node in G.nodes():
        G.nodes[node]['feature'] = random.choice(['a', 'b'])

    # 为每条边添加特征
    for u, v in G.edges():
        G.edges[u, v]['feature'] = random.choice(['a', 'c'])

    return G


def visualize_topology(G, title="random_topology"):
    """
    可视化网络拓扑
    :param G: 图对象
    :param title: 图的标题
    """
    pos = nx.spring_layout(G)  # 使用 spring 布局
    nx.draw_networkx_nodes(G, pos, node_color='lightblue', node_size=500)
    nx.draw_networkx_edges(G, pos, edge_color='gray')
    nx.draw_networkx_labels(G, pos, font_size=10, font_color='black')
    plt.title(title)
    plt.show(block=True)


def export_topology_data(G, file_name="graph_data.json"):
    """
    导出拓扑数据到 JSON 文件
    :param G: 图对象
    :param file_name: 输出文件名
    """
    data = json_graph.node_link_data(G)
    with open(file_name, 'w') as f:
        json.dump(data, f, indent=4)
    print(f"拓扑数据已导出到 {file_name}")


def main():
    # 参数设置
    num_nodes = 50  # 节点数量
    edge_probability = 0.1  # 每对节点之间存在边的概率

    # 生成随机拓扑
    G = generate_random_topology(num_nodes, edge_probability)

    # 可视化拓扑
    visualize_topology(G)

    # 导出数据
    export_topology_data(G, "graph_data.json")


if __name__ == "__main__":
    main()
